from argparse import ArgumentParser
import wave
import numpy as np
import os
import os.path as path
import matplotlib.pyplot as plt
from random import shuffle
import math
def read_wav_data(wavPath):
    '''
    读取一个wav文件，返回声音信号的时域谱矩阵和播放时间
    '''
    wav = wave.open(wavPath,"rb") # 打开一个wav格式的声音文件流
    num_frame = wav.getnframes() # 获取帧数
    num_channel=wav.getnchannels() # 获取声道数
    assert(num_channel == 1)
    framerate=wav.getframerate() # 获取帧速率
    num_sample_width=wav.getsampwidth() # 获取实例的比特宽度，即每一帧的字节数
    #print('samples',num_frame,'  channels',num_channel,'  sample_rate',framerate, 'bits',
    #8 * num_sample_width)
    str_data = wav.readframes(num_frame) # 读取全部的帧
    wav.close() # 关闭流
    wave_data = np.frombuffer(str_data, dtype = np.short) # 将声音文件数据转换为数组矩阵形式
    # print( wave_data.shape)
    wave_data.shape = -1, num_channel # 按照声道数将数组整形，单声道时候是一列数组，双声道时候是两列的矩阵
    # print( wave_data.shape)
    wave_data = np.reshape(wave_data, (-1))
    # print( wave_data.shape)
    return wave_data, framerate  

def wav_show(wave_data, fs): # 显示出来声音波形
    time = np.arange(0, len(wave_data)) * (1.0/fs)  # 计算声音的播放时间，单位为秒
    # 画声音波形
    plt.plot(time, wave_data)  
    plt.show()  


class KwSpliter():
    def __init__(self, workDir='./', itvDurMS=50, quietPowerDbThre=-30, kwDurMS=1000, 
        leadingQuietMinMS=100, leadingQuietMaxMS=300, outRootDir = './kws_dataset_generated'):
        self.dctKwCnt:dict = {}
        
        workDir = workDir.replace('\\', '/')
        while workDir[-1] == '/':
            workDir = workDir[:-1]
        self.workDir = workDir
        self.itvDurMS = itvDurMS
        self.quietPowerDbThre = quietPowerDbThre
        self.kwDurMS = kwDurMS
        self.leadingQuietMinMS = leadingQuietMinMS
        self.leadingQuietMaxMS = leadingQuietMaxMS
        self.outRootDir = outRootDir

    def SplitAllWavFiles(self, workDir=None):
        if workDir is None:
            workDir = self.workDir
        outDirAbs = path.abspath(self.outRootDir).lower()
        for root, dirs, files in os.walk(workDir, topdown=False):
            rootAbs = path.abspath(root).lower()
            if rootAbs.find(outDirAbs) >= 0:
                continue
            for name in files:
                if name[-4:].lower() != '.wav':
                    continue
                wavPath = path.join(root, name)
                wavPath = wavPath.replace('\\', '/')                
                print('spliting %s' % (wavPath))
                self.SplitKW(wavPath, self.itvDurMS, self.quietPowerDbThre, self.kwDurMS,
                    self.leadingQuietMinMS, self.leadingQuietMaxMS, self.outRootDir)
                

    def SplitKW(self, wavPath:str, itvDurMS=50, quietPowerDbThre=-30, kwDurMS=1000, leadingQuietMinMS=100, 
        leadingQuietMaxMS=300, outRootDir = './kws_dataset_generated'):
        if kwDurMS % itvDurMS != 0:
            raise ValueError('keyword duration (ms) must be divisible by interval duration (ms)')
        assert(kwDurMS % itvDurMS == 0)
        ary:np.ndarray = None
        ary, sampRate = read_wav_data(wavPath)

        # strip remainder
        ptsPerItv = int(np.round(sampRate * itvDurMS / 1000.0))
        remainder = len(ary) % ptsPerItv
        if remainder != 0:
            # rocky: only strip remainder when it is not 0
            ary = ary[:-remainder]

        # 求出以itvDurMS为单位的平均能量，以dB为单位
        aryNormed = ary / 32768.0
        aryPower:np.ndarray = aryNormed**2
        aryPower = aryPower.reshape((-1, ptsPerItv))
        aryPowerMean =  aryPower.mean(axis=1)
        aryPowerMean /= aryPowerMean.max()
        aryPowerMean += 1E-8
        aryPowerMean = np.log10(aryPowerMean) * 20

        # aryPowerMean = np.math.log10(aryPowerMean)
        # plt.plot(np.arange(len(aryPowerMean)), aryPowerMean)
        # plt.show()

        # 使用一个状态机来分割口令数据。交替分析安静区和数据区
        WANT_QUIET = 0
        WANT_DATA = 1
        status = WANT_QUIET

        # 话音中可能有字间隔, 设置字间隔的最大允许时长为250ms
        dataToQuietItvCnt = np.math.ceil(250 / itvDurMS)
        lstDataItvRanges = []

        status = WANT_QUIET
        # 先至少检测到连接leadingQuietItvNeeded个安静时隙才检测数据
        leadingQuietItvNeeded = 3
        quietItvToWantData = leadingQuietItvNeeded
        leadingQuietItvCnt = 0
        
        dataStartItv = 0
        for i in range(len(aryPowerMean)):
            if status == WANT_QUIET:
                if aryPowerMean[i] <= quietPowerDbThre:
                    leadingQuietItvCnt += 1
                    quietItvToWantData -= 1
                    if quietItvToWantData == 0:
                        status = WANT_DATA
                        dataItvCnt = 0
                        dataToQuietItvRemaining = dataToQuietItvCnt
                else:
                    # 在需要安静的区域检测到数据, 从头开始检测
                    quietItvToWantData = leadingQuietItvNeeded
                    leadingQuietItvCnt = 0                
            elif status == WANT_DATA:
                if aryPowerMean[i] <= quietPowerDbThre:
                    # 在需要数据的区域检测到安静
                    if dataItvCnt > 0:
                        # 在这之前已经有数据了, 检测是否是数据区中的停顿
                        dataToQuietItvRemaining -= 1
                        if dataToQuietItvRemaining == 0:
                            # 数据区结束, 记录本次数据区的起始和终止
                            lstDataItvRanges.append((dataStartItv, i - dataToQuietItvCnt + 1))
                            # 复位状态机
                            status = WANT_QUIET
                            leadingQuietItvNeeded = 3
                            quietItvToWantData = leadingQuietItvNeeded
                            leadingQuietItvCnt = 0
                        else:
                            # 遇到了数据区中间的短暂安静, 仍然算作是数据区
                            dataItvCnt += 1
                    else:
                        # 在这之前还没收到数据, 继续算作是前导安静区
                        leadingQuietItvCnt += 1
                else:
                    dataToQuietItvRemaining = dataToQuietItvCnt
                    if dataItvCnt == 0:
                        # 记录数据区起始的itv
                        dataStartItv = i
                    dataItvCnt += 1

        # 计算各个口令的起始时隙号
        lstKwItvs = []
        kwDurItvCnt = kwDurMS // itvDurMS
        for (i, item) in enumerate(lstDataItvRanges):

            datDurItvCnt = item[1] - item[0]
            quietItvCnt = kwDurItvCnt - datDurItvCnt
            if quietItvCnt <= 0:
                leadingQuietItvCnt = 0
                if quietItvCnt < 0:
                    print('warning: one keyword audio is longer than %dms!' % (kwDurMS))
            else:
                # 为突出数据集的多样性, 随机在允许的范围设置前导安静时隙数
                leadingQuietItvCnt = np.random.randint(0, quietItvCnt)
                if leadingQuietItvCnt > leadingQuietMaxMS // itvDurMS:
                    leadingQuietItvCnt = leadingQuietMaxMS // itvDurMS
                if leadingQuietItvCnt < leadingQuietMinMS // itvDurMS:
                    leadingQuietItvCnt = leadingQuietMinMS // itvDurMS

                if leadingQuietItvCnt + item[0] >= item[1]:
                    leadingQuietItvCnt -= item[0] + leadingQuietItvCnt - item[1] + 1

                if i == 0:
                    if item[0] < leadingQuietItvCnt:
                        print('Warning: Leading quiet period is too short')
                        leadingQuietItvCnt = item[0]
                else:
                    if item[0] - lstDataItvRanges[i - 1][1] < leadingQuietItvCnt:
                        print('Warning: Consequtive keywords are Too close')
                        leadingQuietItvCnt =  item[0] - lstDataItvRanges[i - 1][1]

            kwItv0 = item[0] - leadingQuietItvCnt
            if kwItv0 + kwDurItvCnt < len(aryPowerMean):
                lstKwItvs.append(kwItv0)


        wavFileName:str = path.split(wavPath)[-1]
        wavMainName:str = path.splitext(wavFileName)[0]
        keywordName:str = wavMainName.split('_')[0]

        if not keywordName in self.dctKwCnt.keys():
            self.dctKwCnt[keywordName] = 0

        outRootDir = outRootDir.replace('\\', '/')
        while outRootDir[-1] == '/':
            outRootDir = outRootDir[:-1]
        
        kwRootPath = path.abspath('%s/%s' % (outRootDir, keywordName)).replace('\\', '/')
        
        if not path.exists(kwRootPath):
            os.makedirs(kwRootPath)
        for (i, itv) in enumerate(lstKwItvs):
            kwAudData = ary[itv * ptsPerItv : (itv + kwDurItvCnt) * ptsPerItv]

            t=np.arange(0,kwDurMS,1000/sampRate) #定义时间点
            kwFilePath = '%s/%s_%03d.wav' % \
                (kwRootPath, wavMainName, self.dctKwCnt[keywordName])
            self.dctKwCnt[keywordName] += 1
            f:wave.Wave_write =wave.open(kwFilePath,'wb') #打开wav声音文件
            f.setnchannels(1) #定义声道数为1
            f.setsampwidth(2) #定义量化位数 2*8=16
            f.setframerate(sampRate) #定义采样率
            f.writeframes(kwAudData.tobytes()) #写入wav声音文件
            f.close()


        return lstKwItvs

if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('-d', '--dir', type=str, default = '.',
        help='working dir that contain .wav files to split;  default to current dir')

    parser.add_argument(
        '-l','--length', type=int, default=1000,
        help='duration of a keyword in ms;  default = 1000(ms)')

    parser.add_argument(
        '-i','--interval', type=int, default=50,
        help='parsing interval ms, must be divisible by length;  default = 50(ms)')

    parser.add_argument(
        '-q','--quiet_db', type=int, default=-30,
        help='max power in dB of quiet;  default = -30(dB)')

    parser.add_argument(
        '--min_lead', type=int, default=50,
        help='min leading quiet duration in ms;  default = 100(ms)')

    parser.add_argument(
        '--max_lead', type=int, default=200,
        help='max leading quiet duration in ms;  default = 300(ms)')

    args,unparsed = parser.parse_known_args()

    if not path.exists(args.dir):
        raise ValueError('working path \'%s\' does not exist' % (args.dir))

    splitObj = KwSpliter(args.dir, itvDurMS=args.interval, kwDurMS=args.length, 
        quietPowerDbThre=args.quiet_db, 
        leadingQuietMinMS=args.min_lead, leadingQuietMaxMS=args.max_lead)
    splitObj.SplitAllWavFiles()
    print('done')

    